#
import unittest
from nanods.nanods_data_curator import NanodsDataCurator

class TNanodsDataCurator(unittest.TestCase):
    # python -m unittest -v uts.nanods.t_nanods_data_curator.TNanodsDataCurator.test_load_tokenizer
    def test_load_tokenizer(self):
        tokenizer = NanodsDataCurator.load_tokenizer()
        print(f'加载的tokenizer词表大小: {len(tokenizer)}')

    # python -m unittest -v uts.nanods.t_nanods_data_curator.TNanodsDataCurator.test_preview_dataset
    def test_preview_dataset(self):
        # 指定文件路径和需要展示的行数
        file_path = './datasets/mobvoi_seq_monkey_general_open_corpus.jsonl'
        NanodsDataCurator.preview_dataset(file_path, num_lines=5)

    # python -m unittest -v uts.nanods.t_nanods_data_curator.TNanodsDataCurator.test_get_total_lines
    def test_get_total_lines(self):
        # 指定文件路径和需要展示的行数
        file_path = './datasets/mobvoi_seq_monkey_general_open_corpus.jsonl'
        nums = NanodsDataCurator.get_total_lines(file_path=file_path)
        print(f'总行数为：{nums};')

    # python -m unittest -v uts.nanods.t_nanods_data_curator.TNanodsDataCurator.test_check_jsonl_format
    def test_check_jsonl_format(self):
        # 指定文件路径和需要展示的行数
        file_path = './datasets/mobvoi_seq_monkey_general_open_corpus.jsonl'
        valid_lines, invalid_lines, invalid_lns = NanodsDataCurator.check_jsonl_format(file_path=file_path)
        print(f"检查完成，文件中共有 {valid_lines} 行有效的 JSON 数据，{invalid_lines} 行无效的 JSON 数据。")
        print(f'无效行：\n{invalid_lns};')

    # python -m unittest -v uts.nanods.t_nanods_data_curator.TNanodsDataCurator.test_remove_invalid_lines
    def test_remove_invalid_lines(self):
        file_path = './datasets/mobvoi_seq_monkey_general_open_corpus.jsonl'
        valid_lines, invalid_lines, invalid_lns = NanodsDataCurator.check_jsonl_format(file_path=file_path)
        # 使用该函数删除第 9598787 行并保存为新文件
        raw_ds_fn = './datasets/mobvoi_seq_monkey_general_open_corpus.jsonl'
        ds_fn = './datasets/mobvoi_seq_monkey_general_open_corpus_cleaned.jsonl'
        NanodsDataCurator.remove_invalid_lines(raw_ds_fn,
                            ds_fn, 
                            invalid_lns=invalid_lns)
        
    # python -m unittest -v uts.nanods.t_nanods_data_curator.TNanodsDataCurator.test_pretrain_process
    def test_pretrain_process(self):
        NanodsDataCurator.pretrain_process()

    # python -m unittest -v uts.nanods.t_nanods_data_curator.TNanodsDataCurator.test_check_ds
    def test_check_ds(self):
        NanodsDataCurator.check_ds()